package com.cwj.algorithm.tree;

import java.util.LinkedList;
import java.util.Queue;

/**
 * @author chenwujie
 * @date 2020-12-17 14:51
 */
public class RedBlackTree<K extends Comparable<? super K>> {
    private TreeNode<K> root;
    private int size;

    public void put(K k){
        TreeNode<K> newNode = new TreeNode<>(k, null, null);
        if (root == null) {
            newNode.parent = null;
            newNode.red = false;
            root = newNode;
        }else{
            put(root, newNode);
            balanceInsertion(newNode);
        }
        size++;
    }
    private void put(TreeNode<K> node, TreeNode<K> x){
        if(node.k.compareTo(x.k) > 0){
            if(node.left == null){
                x.parent = node;
                node.left = x;
            }else{
                put(node.left, x);
            }
        }else{
            if(node.right == null){
                x.parent = node;
                node.right = x;
            }else{
                put(node.right, x);
            }
        }
    }

    @Override
    public String toString() {
        Queue<RedBlackTree.TreeNode<K>> queue = new LinkedList<>();
        Queue<RedBlackTree.TreeNode<K>> temp = new LinkedList<>();
        if(root == null){
            return "null";
        }
        temp.add(root);
        while(!temp.isEmpty()){
            RedBlackTree.TreeNode<K> poll = temp.poll();
            queue.add(poll);
            if(poll.left != null){
                temp.add(poll.left);
            }
            if(poll.right != null){
                temp.add(poll.right);
            }
        }
        StringBuilder stringBuilder = new StringBuilder();
        queue.forEach(e -> stringBuilder.append(e.toString()).append(" "));
        return stringBuilder.toString();
    }

    public TreeNode<K> getNode(K k){
        return getNode(root, k);
    }

    private TreeNode<K> getNode(TreeNode<K> node, K k){
        if(node == null){
            return null;
        }
        if(node.k.compareTo(k) == 0){
            return node;
        }
        if(node.k.compareTo(k) > 0){
            return getNode(node.left, k);
        }else{
            return getNode(node.right, k);
        }
    }

    private TreeNode<K> balanceInsertion(TreeNode<K> x){
        for (TreeNode<K> xp, xpp, xppl, xppr;;) {
            if ((xp = x.parent) == null) {
                x.red = false;
                return x;
            }
            else if (!xp.red || (xpp = xp.parent) == null) {
                return root;
            }
            if (xp == (xppl = xpp.left)) {
                if ((xppr = xpp.right) != null && xppr.red) {
                    xppr.red = false;
                    xp.red = false;
                    xpp.red = true;
                    x = xpp;
                }
                else {
                    if (x == xp.right) {
                        root = rotateLeft(x = xp);
                        xpp = (xp = x.parent) == null ? null : xp.parent;
                    }
                    if (xp != null) {
                        xp.red = false;
                        if (xpp != null) {
                            xpp.red = true;
                            root = rotateRight(xpp);
                        }
                    }
                }
            }
            else {
                if (xppl != null && xppl.red) {
                    xppl.red = false;
                    xp.red = false;
                    xpp.red = true;
                    x = xpp;
                }
                else {
                    if (x == xp.left) {
                        root = rotateRight(x = xp);
                        xpp = (xp = x.parent) == null ? null : xp.parent;
                    }
                    if (xp != null) {
                        xp.red = false;
                        if (xpp != null) {
                            xpp.red = true;
                            root = rotateLeft(xpp);
                        }
                    }
                }
            }
        }
    }

    private TreeNode<K> balanceDeletion(TreeNode<K> x){
        for (TreeNode<K> xp, xpl, xpr;;) {
            if (x == null || x == root) {
                return root;
            } else if ((xp = x.parent) == null) {
                x.red = false;
                return x;
            }
            else if (x.red) {
                x.red = false;
                return root;
            }
            else if ((xpl = xp.left) == x) {
                if ((xpr = xp.right) != null && xpr.red) {
                    xpr.red = false;
                    xp.red = true;
                    root = rotateLeft(xp);
                    xpr = (xp = x.parent) == null ? null : xp.right;
                }
                if (xpr == null) {
                    x = xp;
                } else {
                    TreeNode<K> sl = xpr.left, sr = xpr.right;
                    if ((sr == null || !sr.red) &&
                            (sl == null || !sl.red)) {
                        xpr.red = true;
                        x = xp;
                    }
                    else {
                        if (sr == null || !sr.red) {
                            if (sl != null) {
                                sl.red = false;
                            }
                            xpr.red = true;
                            root = rotateRight(xpr);
                            xpr = (xp = x.parent) == null ?
                                    null : xp.right;
                        }
                        if (xpr != null) {
                            xpr.red = (xp == null) ? false : xp.red;
                            if ((sr = xpr.right) != null) {
                                sr.red = false;
                            }
                        }
                        if (xp != null) {
                            xp.red = false;
                            root = rotateLeft(xp);
                        }
                        x = root;
                    }
                }
            }
            else { // symmetric
                if (xpl != null && xpl.red) {
                    xpl.red = false;
                    xp.red = true;
                    root = rotateRight(xp);
                    xpl = (xp = x.parent) == null ? null : xp.left;
                }
                if (xpl == null) {
                    x = xp;
                } else {
                    TreeNode<K> sl = xpl.left, sr = xpl.right;
                    if ((sl == null || !sl.red) &&
                            (sr == null || !sr.red)) {
                        xpl.red = true;
                        x = xp;
                    }
                    else {
                        if (sl == null || !sl.red) {
                            if (sr != null) {
                                sr.red = false;
                            }
                            xpl.red = true;
                            root = rotateLeft(xpl);
                            xpl = (xp = x.parent) == null ?
                                    null : xp.left;
                        }
                        if (xpl != null) {
                            xpl.red = (xp == null) ? false : xp.red;
                            if ((sl = xpl.left) != null) {
                                sl.red = false;
                            }
                        }
                        if (xp != null) {
                            xp.red = false;
                            root = rotateRight(xp);
                        }
                        x = root;
                    }
                }
            }
        }
    }

    private TreeNode<K> rotateLeft(TreeNode<K> p){
        TreeNode<K> r, pp, rl;
        if (p != null && (r = p.right) != null) {
            if ((rl = p.right = r.left) != null) {
                rl.parent = p;
            }
            if ((pp = r.parent = p.parent) == null) {
                (root = r).red = false;
            } else if (pp.left == p) {
                pp.left = r;
            } else {
                pp.right = r;
            }
            r.left = p;
            p.parent = r;
        }
        return root;
    }

    private TreeNode<K> rotateRight(TreeNode<K> p){
        TreeNode<K> l, pp, lr;
        if (p != null && (l = p.left) != null) {
            if ((lr = p.left = l.right) != null) {
                lr.parent = p;
            }
            if ((pp = l.parent = p.parent) == null) {
                (root = l).red = false;
            } else if (pp.right == p) {
                pp.right = l;
            } else {
                pp.left = l;
            }
            l.right = p;
            p.parent = l;
        }
        return root;
    }

    static class TreeNode<K>{
        K k;
        TreeNode<K> parent;
        TreeNode<K> left;
        TreeNode<K> right;
        boolean red;

        TreeNode(K k, TreeNode<K> left, TreeNode<K> right){
            this.k = k;
            this.left = left;
            this.right = right;
            this.red = true;
        }

        @Override
        public String toString() {
            String color = red ? "红色" : "黑色";
            return k.toString() + ": " + color;
        }
    }
}

